package task_process

import (
	"context"
	"fmt"
	"log"
	"time"

	"github.com/gogf/gf/container/gmap"
	"github.com/gogf/gf/frame/g"
	"github.com/hibiken/asynq"
)

const (
	TORCHSR_SERVER_TASK = "torchsr_server_task"
)

func init() {
	AddMuxPointer(TORCHSR_SERVER_TASK, ProcessImageTo3DModel, 0, 7200)
}

type Taskoption struct {
	RetryCount int64 //重试次数
	TimeOut    int64 //最大运行秒数
}

type HandlerMap map[string]func(c context.Context, t *asynq.Task) error

var optMap *gmap.Map

var allMap HandlerMap = nil

var schAr []*asynq.Scheduler = nil

var ScheduleMap map[string]string = nil

var avaliableGPUMap *gmap.IntIntMap

func AddScheduleTask(corn, taskName string) {
	if ScheduleMap == nil {
		ScheduleMap = make(map[string]string)
	}
	ScheduleMap[taskName] = corn
}

func AddMuxPointer(tag string, profunc func(c context.Context, t *asynq.Task) error, retryTimes int64, timeOut int64) {

	if allMap == nil {
		allMap = make(map[string]func(c context.Context, t *asynq.Task) error)
	}

	if optMap == nil {
		optMap = gmap.New(true)
	}

	topt := new(Taskoption)
	topt.RetryCount = retryTimes
	topt.TimeOut = timeOut
	optMap.Set(tag, topt)
	allMap[tag] = profunc
}

func StartServer() *asynq.Server {
	mux := asynq.NewServeMux()
	for k, v := range allMap {
		mux.HandleFunc(k, v)
	}

	avaliableGPUIds := g.Cfg().GetInts("avaliable_gpu")
	concurrencyCount := len(avaliableGPUIds)

	if avaliableGPUMap == nil {
		avaliableGPUMap = gmap.NewIntIntMap(true)
		for _, element := range avaliableGPUIds {
			avaliableGPUMap.Set(element, 1)
		}
	}

	taskRetryDelay := 5
	r := GetRedisCon()
	srv := asynq.NewServer(r, asynq.Config{
		// Specify how many concurrent workers to use
		Concurrency: concurrencyCount,
		// Optionally specify multiple queues with different priority.
		Queues: map[string]int{
			"critical": 6,
			"default":  3,
			"low":      1,
		},
		// See the godoc for other configuration options
		RetryDelayFunc: func(n int, e error, t *asynq.Task) time.Duration {
			return time.Second * time.Duration(taskRetryDelay)
		},
	})

	// mux maps a type to a handler
	//mux := asynq.NewServeMux()
	//mux.HandleFunc(EmailDelivery,HandleEmailDeliveryTask)
	//mux.Handle(ImageProcessing, NewImageProcessor())
	// 添加自己的任务处理

	if err := srv.Start(mux); err != nil {
		log.Fatalf("could not run server: %v", err)
	}
	startPeriodicClearTask()
	fmt.Println("StartServer")
	return srv
}

func startPeriodicClearTask() {
	r := GetRedisCon()

	schAr := make([]*asynq.Scheduler, 0)

	for k, v := range ScheduleMap {
		taskName := k
		taskCron := v

		scheduler := asynq.NewScheduler(r, nil)
		task := asynq.NewTask(taskName, nil)
		_, err := scheduler.Register(taskCron, task)
		if err != nil {
			log.Fatalf("scheduler.Register fail: %v", err)
		}
		err = scheduler.Start()
		if err != nil {
			log.Fatalf("scheduler.Register fail: %v", err)
		}
		schAr = append(schAr, scheduler)
	}
}

func StopServer(s *asynq.Server) {
	for _, v := range schAr {
		sc := v
		sc.Stop()
	}
	s.Stop()
}
